import pyaudio
import webrtcvad
import numpy as np

# 参数配置
AUDIO_RATE = 16000       # 采样率：16000（支持8000, 16000, 32000或48000）
CHUNK_SIZE = 480         # 每块大小（30ms，保证为10/20/30ms的倍数）
VAD_MODE = 1             # VAD模式（0-3，数值越小越保守）

# 初始化VAD
vad = webrtcvad.Vad(VAD_MODE)

class SpeechDetector:
    def __init__(self, calibration_seconds=2, chunk_duration_ms=30):
        """
        calibration_seconds: 校准背景噪音所需时间（秒）
        chunk_duration_ms: 每块时长（毫秒）
        """
        self.calibration_seconds = calibration_seconds
        self.chunk_duration_ms = chunk_duration_ms
        self.calibrated = False
        self.amplitude_threshold = None

        # 连续帧判决参数（降低短时噪音误判）
        self.speech_state = False         # 当前状态：True为语音，False为无语音
        self.consecutive_speech = 0       # 连续语音帧计数
        self.consecutive_silence = 0      # 连续静音帧计数
        self.required_speech_frames = 3   # 连续3帧语音后确认进入语音状态（约90ms）
        self.required_silence_frames = 5*5  # 750ms 连续5帧静音后退出语音状态（约150ms）

    def calibrate(self, stream):
        """
        校准背景噪音：录制固定时长音频，计算平均幅值与标准差，从而设置自适应阈值
        """
        print("开始校准背景噪音，请保持安静...")
        amplitudes = []
        num_frames = int(self.calibration_seconds * (1000 / self.chunk_duration_ms))
        for _ in range(num_frames):
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
            amplitudes.append(np.abs(audio_data).mean())
        mean_noise = np.mean(amplitudes)
        std_noise = np.std(amplitudes)
        # 阈值设置为均值加2倍标准差
        self.amplitude_threshold = mean_noise + 2 * std_noise
        print(f"校准完成：噪音均值={mean_noise:.2f}，标准差={std_noise:.2f}，设置阈值={self.amplitude_threshold:.2f}")
        self.calibrated = True

    def analyze_spectrum(self, audio_chunk):
        """
        通过频谱分析检测语音特性：
        1. 对音频块应用汉宁窗后计算 FFT
        2. 统计局部峰值数量（峰值必须超过均值的1.5倍）
        3. 当峰值数量大于等于设定阈值时，认为该块具有语音特征
        """
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        if len(audio_data) == 0:
            return False

        # 使用汉宁窗减少FFT泄露
        window = np.hanning(len(audio_data))
        windowed_data = audio_data * window

        # FFT计算得到频谱（只需正频率部分）
        spectrum = np.abs(np.fft.rfft(windowed_data))

        # 使用均值作为参考，统计超过均值1.5倍的局部峰值数量
        spectral_mean = np.mean(spectrum)
        peak_count = 0
        for i in range(1, len(spectrum) - 1):
            if spectrum[i] > spectrum[i - 1] and spectrum[i] > spectrum[i + 1] \
               and spectrum[i] > spectral_mean * 1.5:
                peak_count += 1

        # 参数阈值：这里设定至少需要3个峰值（可根据实际情况调整）
        return peak_count >= 3

    def is_speech(self, audio_chunk):
        """
        综合判断：先通过能量预处理（阈值）过滤低能量数据，再利用VAD和频谱分析判断语音。
        两者结合能有效降低噪音导致的误判。
        """
        # 使用背景校准得到的阈值（若未校准则取较低默认值）
        amplitude_threshold = self.amplitude_threshold if self.amplitude_threshold is not None else 500

        # 能量检测
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        amplitude = np.abs(audio_data).mean()
        if amplitude < amplitude_threshold:
            return False

        # VAD检测
        vad_result = vad.is_speech(audio_chunk, AUDIO_RATE)
        # 频谱检测
        spectral_result = self.analyze_spectrum(audio_chunk)

        # 仅当两种检测均为True时，认为当前块为语音
        return vad_result and spectral_result

    def process_chunk(self, audio_chunk):
        """
        对每一块数据进行处理：综合能量检测、VAD、频谱分析，并采用连续帧策略实现状态平滑。
        """
        is_speech_chunk = self.is_speech(audio_chunk)

        # 连续帧计数策略
        if is_speech_chunk:
            self.consecutive_speech += 1
            self.consecutive_silence = 0
        else:
            self.consecutive_silence += 1
            self.consecutive_speech = 0

        # 状态转换逻辑
        if not self.speech_state and self.consecutive_speech >= self.required_speech_frames:
            self.speech_state = True
            print("Detected Speech")
        elif self.speech_state and self.consecutive_silence >= self.required_silence_frames:
            self.speech_state = False
            print("No speech")

def main():
    p = pyaudio.PyAudio()
    stream = p.open(format=pyaudio.paInt16,
                    channels=1,
                    rate=AUDIO_RATE,
                    input=True,
                    frames_per_buffer=CHUNK_SIZE)
    
    detector = SpeechDetector()
    # 校准背景噪音（建议在程序启动时进行）
    detector.calibrate(stream)
    
    print("开始监听，请开始说话...(按Ctrl+C停止)")
    try:
        while True:
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            detector.process_chunk(audio_chunk)
    except KeyboardInterrupt:
        print("停止监听")
    finally:
        stream.stop_stream()
        stream.close()
        p.terminate()

if __name__ == "__main__":
    main()
